Lambda + EventBridgeでSageMaker StudioのKernelGatewayを自動削除して不要なコストを削減する

貞松政史

2021.05.30

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

データアナリティクス事業本部の貞松です。最近は日々SageMakerと格闘しております。

機械学習のWebベースIDEであるSageMaker Studioは、Jupyter Notebookによる開発時にKernelという単位で管理される環境を利用することで、Notebookを実行する為のインスタンスや各種設定やライブラリが搭載されたコンテナ環境を柔軟に切り替えながら開発を行うことができます。

実行されているKernelはKernelGateWayというタイプのアプリケーションとして、SageMaker Studioの実行ユーザー単位で作成されるのですが、このKernelGatewayが起動している間は、Kernelを実行しているインスタンスタイプに応じて時間単位で課金が発生します。

また、KernelGatewayは、SageMaker Studioの画面を閉じたり、SageMaker Studioの本体(Jupyter Server)をシャットダウンしても停止しない為、ユーザーが明示的にKernelのインスタンスをシャットダウンする、あるいは該当のKernelGatewayタイプのアプリケーションを削除する必要があります。

AWSマネジメントコンソールの画面やSageMaker Studioの画面から手動で削除しても良いのですが、毎回手動で削除するのは手間がかかりますし、何より削除忘れによる無駄な課金が発生してしまう可能性があります。

そこで今回は、AWS SDK経由でSageMaker StudioのKernelGatewayを削除する処理をLambda関数として実装し、その関数をEventBridgeからスケジュール実行する仕組みを作成します。

SageMaker StudioのKernelGatewayを削除するLambda関数を作成する

開発環境

今回作成したLambda関数のランタイムはPython3.8を使用しています。

作成した関数の処理仕様

今回作成する処理の仕様は以下の通りです。

SageMaker StudioのDomainIdを環境変数に設定して、該当のDomainIdを持つSageMaker StudioのAppのリストを取得する。
取得したAppのリストからAppType=KernelGateweyかつStatus=InServiceのAppのUserProfileNameとAppNameを取得して別のリストに詰め直す。
詰め直したリストに含まれるApp(KernelGateway)を全て削除する。

実行に必要なIAMロールの権限設定

上記の仕様に沿って処理を実行する為に、Lambdaの実行ロールに以下の権限を追加します。

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": [
                "sagemaker:ListApps",
                "sagemaker:DeleteApp"
            ],
            "Resource": "*"
        }
    ]
}

環境変数の設定

Lambda関数の環境変数としてDOMAIN_IDを設定します。

SageMaker StudioのDomainIdはAWSマネジメントコンソールのSageMaker画面上ではStudio IDとして表示されているので、この値を設定します。

Pythonコードの作成

AWS SDK for Python (Boto3)を使用して、SagaMaker StudioからAppのリストを取得し、KernelGatewayを削除する処理を実装します。
今回使用するBoto3のSageMakerクライアントの仕様については下記URLのページをご参照ください。

Boto3 - SageMaker.Client.list_apps

Boto3 - SageMaker.Client.delete_app

以下、作成したLambda関数のPythonコードです。
前述で設定した環境変数を先頭で読み込んで使用しています。

# sagemaker_studio_delete_apps.ｐｙ

import boto3
import json
import os

domain_id = os.environ['DOMAIN_ID']

def delete_apps(event, context):
    client = boto3.Session().client('sagemaker')

    # 指定されたのDomainIdのSageMaker StudioからAppのリストを取得
    response = client.list_apps(
        DomainIdEquals=domain_id
    )
    
    # 取得したAppのリストから
    # AppTypeがKernelGateweyかつStatusがInServiceのAppの
    # UserProfileNameとAppNameを取得してリストに詰め直す
    app_list = []
    apps = response['Apps']
    for app in apps:
        if app['AppType'] == 'KernelGateway' and app['Status'] == 'InService':
            app_list.append((app['UserProfileName'], app['AppName']))
    
    # 詰め直したリストに含まれるApp(KernelGateway)を削除する
    for app in app_list:
        client.delete_app(
            DomainId = domain_id,
            UserProfileName = app[0],
            AppType = 'KernelGateway',
            AppName = app[1]
        )

    return {
        'statusCode': 200,
        'body': json.dumps('Started deleting apps.')
    }

今回作成したPythonのファイル名はsagemaker_studio_delete_appsで関数名をdelete_appsとしているので、ランタイム設定のハンドラはsagemaker_studio_delete_apps.delete_appsとなっています。ここは適宜作成したファイルの内容に合わせて書き換えてください。